from langchain_community.chat_message_histories import SQLChatMessageHistory
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory, RunnablePassthrough
import gradio
from my_llm import llm

# 1 提示器模版
prompt = ChatPromptTemplate.from_messages([
    ('system', '{system_message}'),
    MessagesPlaceholder(variable_name='chat_history', optional=True),
    # ('placeholder', 'chat_history'),
    ('human', '{input}')
])

chain = prompt | llm


# 2 存储所有会话的聊天记录（内存，关系型数据库或者redis）
# store = {}  # 保存历史聊天记录key：sessionID，保存到数据库则不需要创建


def get_session_history(session_id: str):
    """
    从关系型数据库聊天历史记录中取聊天历史记录
    :param session_id:
    :return:
    """
    return SQLChatMessageHistory(
        session_id=session_id,
        connection_string='sqlite:///chat_history.db',
    )


# 消息类型 SystemMessage、HumanMessage、AIMessage、ToolMessage

# 3 创建带有历史记录功能的处理链

chain_with_message_history = RunnableWithMessageHistory(
    chain,
    get_session_history,
    input_messages_key='input',  # 用户输入的消息的key
    history_messages_key='chat_history'  # 就是历史消息占位符中的variable_name
)

"""
保存和摘要上下文历史记录：保存最近的前2条消息，把之前所有的消息形成摘要
"""


def summarize_message(current_input):
    """
    剪辑和摘要上下文，历史记录
    :param current_input:
    :return:
    """
    session_id = current_input['config']['configurable']['session_id']
    if not session_id:
        raise ValueError("必须通过config参数提供sessionID")

    # 获取当前会话ID的所有历史聊天记录
    chat_history = get_session_history(session_id)
    store_messages = chat_history.messages

    if len(store_messages) <= 2:  # 保留最近2条嘻嘻的阈值
        return {
            "original_messages": store_messages,  # 原始的消息
            "summary": None,  # 不满足两条就不生成摘要
        }

    # 剪辑消息列表
    last_two_messages = store_messages[-2:]
    messages_to_summarize = store_messages[:-2]  # 需要进行摘要的消息列表

    summarization_prompt = ChatPromptTemplate.from_messages([
        ("system", "请将以下对话历史压缩为一条保存关键信息的摘要消息"),
        ("placeholder", "{chat_history}"),
        ("human", "请生成包含上述对话核心内容的摘要，保留重要事实和决策"),
    ])

    summarization_chain = summarization_prompt | llm
    # 生成摘要(AIMessage)
    summary_message = summarization_chain.invoke({'chat_history': messages_to_summarize})
    """
    # 重建历史记录：摘要＋最后２条原始消息
    # chat_history.clear()
    # chat_history.add_message(summary_message)
    # for msg in last_two_messages:
    #     chat_history.add_message(msg)
    """
    # 修改方法，返回结构华的结果（chat_history.clear()）
    return {
        "original_messages": last_two_messages,  # 保留的原始消息（一个是HumanMessage,一个是AIMessage）
        "summary": summary_message,  # 生成的摘要(AIMessage)
    }


# 最终的链
# 1、input，原来的messages_summarized=messages_summarized函数执行后的返回值
# 2、input，原来的，chat_history：messages_summarized['original_messages']
# system_message:messages_summarized['summary']
final_chain = (RunnablePassthrough.assign(messages_summarized=summarize_message) | RunnablePassthrough.assign(
    input=lambda x: x['input'],
    chat_history=lambda x: x['messages_summarized']['original_messages'],
    system_message=lambda
        x: f"你是一个乐于助人的助手，尽你所能回答问题，摘要：{x['messages_summarized']['summary'].content}" if
    x['messages_summarized'].get('summary') else "无摘要") | chain_with_message_history)

# 最终链（核心逻辑：生成摘要→构造系统提示→传递给历史链）
# final_chain = (
#     # 1. 第一步：添加摘要字段（`message_summary`）
#     # `summarize_message`函数返回：{'original_messages': [...], 'summary': '...'}
#     RunnablePassthrough.assign(message_summary=summarize_message)
#     # 2. 第二步：从摘要中提取`chat_history`、生成`system_message`（保留原始`input`）
#     | RunnablePassthrough.assign(
#             # 提取保留的原始历史消息（用于`chain_with_message_history`的上下文）
#             chat_history=lambda x: x['message_summary']['original_messages'],
#             # 生成系统提示（结合摘要，无摘要则用默认值）
#             system_message=lambda x: f"你是一个乐于助人的助手，尽你所能回答问题。对话摘要：{x['message_summary']['summary']}"
#             if x['message_summary'].get('summary')  # 正确使用`get`方法（避免KeyError）
#             else "无对话摘要"
# )
#
#     # 3. 第三步：传递给带有历史记录的链（需`input`、`chat_history`、`system_message`）
#     | chain_with_message_history
# )

# 调用执行链，获得结果
# result1 = final_chain.invoke({'input': '你好，我是张廷玉', 'config': {'configurable': {'session_id': 'user123'}}},
#                              config={'configurable': {'session_id': 'user123'}})
# print(result1)
# # 调用执行链，获得结果
# result2 = final_chain.invoke({'input': '我的名字叫什么', 'config': {'configurable': {'session_id': 'user123'}}},
#                              config={'configurable': {'session_id': 'user123'}})
# print(result2)
# 调用执行链，获得结果
# result3 = final_chain.invoke(
#     {'input': '今天天气真好', 'config': {'configurable': {'session_id': 'user123'}}},
#     config={'configurable': {'session_id': 'user123'}})
# print(result3)


"""   开发聊天机器人web界面  """


# web页面核心函数
def add_message(chat_history, user_message):
    if user_message:
        chat_history.append({'role': 'user', 'content': user_message})
    return chat_history, ''  # 返回空字符串，因为聊天记录已经发送，所以重置文本输入框


def execute_chain(chat_history):
    # 输入就是聊天历史的最后一条记录
    input = chat_history[-1]
    result = final_chain.invoke({'input': input['content'], 'config': {'configurable': {'session_id': 'user123'}}},
                                config={'configurable': {'session_id': 'user123'}})
    chat_history.append({'role': 'assistant', 'content': result.content})
    return chat_history


with gradio.Blocks(title='多模态聊天机器人', theme=gradio.themes.Soft()) as block:
    # 聊天历史记录的组件
    chatbot = gradio.Chatbot(type='messages', height=500, label="聊天机器人")

    # 创建多模态输入框
    chat_input = gradio.MultimodalTextbox(
        interactive=True,  # 可交互
        file_types=['image', '.wav', 'mp4'],
        file_count='multiple',  # 允许多文件上传
        placeholder='请输入信息或者上传文件...',
        show_label=False,
        sources=['microphone', 'upload']  # 允许麦克风和文件上传
    )

    with gradio.Row():
        # 文字输入区域
        with gradio.Column(scale=4):
            user_input = gradio.Textbox(placeholder='请给机器人发送消息..', label="文字输入", max_lines=5)
            submit_btn = gradio.Button('发送', variant='primary')
        with gradio.Column(scale=1):
            audio_input = gradio.Audio(sources=['microphone'], label='语音输入', type='filepath', format='wav')

    # 添加事件处理 - 修复点1：为按钮和文本框添加相同的事件处理
    submit_event = submit_btn.click(fn=add_message, inputs=[chatbot, user_input], outputs=[chatbot, user_input])

    # 添加提交事件
    chat_msg = user_input.submit(fn=add_message, inputs=[chatbot, user_input], outputs=[chatbot, user_input])

    # 执行链
    chat_msg.then(execute_chain, chatbot, chatbot)

if __name__ == '__main__':
    block.launch()
